package com.zijidelu.luframework.utils.concurrent;

import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.ThreadPoolExecutor;

/**
 * 线程池工具类。
 *
 * @author ZIJIDELU
 * @datetime 2025/10/30 12:10
 */
public class ThreadPoolUtil {
    /**
     * 根据任务类型自动计算线程池大小
     *
     * @param ioIntensive 是否为 IO 密集型任务
     * @param executor    线程池任务执行器
     * @return ThreadPoolTaskExecutor
     */
    public static ThreadPoolTaskExecutor calcThreadPoolSize(boolean ioIntensive, ThreadPoolTaskExecutor executor) {
        final int cpuCore = Runtime.getRuntime().availableProcessors();
        // 最大线程数限制为合理值
        final int threadLimit = ioIntensive
                ? Math.min(128, cpuCore * 8)  // IO 密集型任务
                : Math.min(64, cpuCore * 2);  // CPU 密集型任务

        // IO 阻塞系数
        final double blockingCoefficient = 1.5;

        // 队列系数
        final double queueCoefficient = 2.0;

        // 队列最小值 / 最大值
        final int queueMin = 10;
        final int queueMax = 200;

        // 核心线程数，避免 CPU 密集型线程数过大
        final int corePoolSize = ioIntensive
                ? (int) Math.min(cpuCore * (1 + blockingCoefficient), (double) threadLimit / 2)
                : Math.min(cpuCore, threadLimit / 2);

        // 最大线程数
        final int maxPoolSize = ioIntensive
                ? threadLimit
                : Math.min(cpuCore * 2, threadLimit);

        // 队列长度，根据核心线程数计算，并限制在合理范围
        final int queueSize = Math.min(
                Math.max((int) (corePoolSize * queueCoefficient), queueMin),
                queueMax
        );

        // keepAliveSeconds 动态计算
        final int keepAliveSeconds = ioIntensive
                ? Math.max(30, Math.min(120, corePoolSize * 2))
                : Math.max(60, Math.min(300, corePoolSize * 5));

        // 核心线程数
        executor.setCorePoolSize(corePoolSize);
        // 最大线程数
        executor.setMaxPoolSize(maxPoolSize);
        // 线程保持时间
        executor.setKeepAliveSeconds(keepAliveSeconds);
        // 允许核心线程超时
        executor.setAllowCoreThreadTimeOut(true);
        // 队列容量
        executor.setQueueCapacity(queueSize);
        // 拒绝策略
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());

        return executor;
    }

}
